﻿#include "qdiskinfoprovider.h"
#include <QProcess>
#include <QStringList>
#include <QDebug>

QDiskInfoProvider::QDiskInfoProvider(QObject *parent)
    : QObject(parent)
{
}

QList<QDiskInfoProvider::DiskDetails> QDiskInfoProvider::getAllDisksDetails() const {
#ifdef Q_OS_WIN
    return getDisksDetailsWindows();
#elif defined(Q_OS_LINUX)
    return getDisksDetailsLinux();
#else
    return {};
#endif
}

QList<QDiskInfoProvider::DiskDetails> QDiskInfoProvider::getDisksDetailsWindows() const {
    QList<DiskDetails> diskDetailsList;
    QProcess process;

    // Use WMI to get detailed information about each disk.
    process.start("wmic", QStringList() << "diskdrive" << "get" << "/format:list");
    process.waitForFinished();
    QString diskInfo = process.readAllStandardOutput();
    QStringList diskBlocks = diskInfo.split("\r\r\n\r\r\n", Qt::SkipEmptyParts);
    foreach (const QString &block, diskBlocks) {
        DiskDetails details;
        details.device = block.section("\nDeviceID=", 1, 1).section('\n', 0, 0).simplified();
        details.model = block.section("\nModel=", 1, 1).section('\n', 0, 0).simplified();
        details.serialNumber = block.section("\nSerialNumber=", 1, 1).section('\n', 0, 0).simplified();
        details.size = block.section("\nSize=", 1, 1).section('\n', 0, 0).simplified(); // Size is in bytes
        details.interfaceType = block.section("\nInterfaceType=", 1, 1).section('\n', 0, 0).simplified();
        details.partitionCount = block.section("\nPartitions=", 1, 1).section('\n', 0, 0).simplified();

        if (!details.device.isEmpty()) { // Assuming at least the DeviceID field should be present for valid entries.
            diskDetailsList.append(details);
        }
    }

    return diskDetailsList;
}

QList<QDiskInfoProvider::DiskDetails> QDiskInfoProvider::getDisksDetailsLinux() const {
    QList<DiskDetails> diskDetailsList;
    QProcess lsblkProcess;
    lsblkProcess.start("lsblk", QStringList() << "-bdo" << "NAME,MODEL,SERIAL,SIZE,TRAN,MAJ:MIN");
    lsblkProcess.waitForFinished();
    QString lsblkOutput = lsblkProcess.readAllStandardOutput();

    QProcess smartctlProcess;
    smartctlProcess.start("smartctl", QStringList() << "--scan-open");
    smartctlProcess.waitForFinished();
    QString smartctlOutput = smartctlProcess.readAllStandardOutput();

    QStringList lines = lsblkOutput.split('\n');
    foreach (const QString &line, lines) {
        if (line.contains("disk")) {
            QStringList parts = line.trimmed().split(' ', Qt::SkipEmptyParts);
            if(parts.length() > 4) {
                DiskDetails details;
                details.device = parts[0]; // NAME
                details.model = parts[1]; // MODEL
                details.serialNumber = ""; // SERIAL might not be directly available through lsblk
                details.size = parts[3]; // SIZE in bytes
                details.interfaceType = parts[4]; // TRAN (Transport type, e.g., sata, nvme)

                // For Serial Number and more accurate Model/Size info, use smartctl on each device
                smartctlProcess.start("smartctl", QStringList() << "-i" << ("/dev/" + details.device));
                smartctlProcess.waitForFinished();
                QString smartctlDeviceOutput = smartctlProcess.readAllStandardOutput();
                details.serialNumber = smartctlDeviceOutput.section("Serial Number:", 1, 1).section('\n', 0, 0).simplified();
                details.model = smartctlDeviceOutput.section("Device Model:", 1, 1).section('\n', 0, 0).simplified();
                details.size = smartctlDeviceOutput.section("User Capacity:", 1, 1).section(',', 0, 0).simplified();

                // Count partitions using lsblk again
                QProcess countPartitionsProcess;
                countPartitionsProcess.start("lsblk", QStringList() << "-ndo" << "TYPE,NAME" << ("/dev/" + details.device));
                countPartitionsProcess.waitForFinished();
                QString countPartitionsOutput = countPartitionsProcess.readAllStandardOutput();
                details.partitionCount = QString::number(countPartitionsOutput.count("part"));

                diskDetailsList.append(details);
            }
        }
    }

    return diskDetailsList;
}
